"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/6 8:55
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    align.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    align_VER0.1

------------------------------------------------------------------
# @Description:

------------------------------------------------------------------
"""
import numpy as np


class RMSLAlign(object):
    def __init__(self,
                 src_w: int, src_h: int,
                 dst_w: int, dst_h: int,
                 extend_w_pixels=256,
                 extend_h_pixels=256,
                 rotate_angle=0):
        self.src_w = src_w
        self.src_h = src_h
        self.dst_w = dst_w
        self.dst_h = dst_h
        self.rotate_angle = rotate_angle
        self.extend_w_pixel = extend_w_pixels
        self.extend_h_pixel = extend_h_pixels

        self.extend_buf_w = dst_w + 2 * self.extend_w_pixel
        self.extend_buf_h = dst_h + 2 * self.extend_h_pixel

        self.lut = np.zeros((self.src_h, self.src_w, 3), dtype=np.float32)

        self.ir_extend = np.zeros((self.extend_buf_h, self.extend_buf_w), dtype=np.uint8)
        self.dp_extend = np.zeros((self.extend_buf_h, self.extend_buf_w), dtype=np.uint16)

    def create_lut(self, r_coeff: dict):
        if len(r_coeff) == 0 or r_coeff is None:
            return

        cols, rows = np.meshgrid(np.arange(self.src_w), np.arange(self.src_h))
        x = r_coeff['r0'] * cols + r_coeff['r1'] * rows + r_coeff['r2']
        y = r_coeff['r3'] * cols + r_coeff['r4'] * rows + r_coeff['r5']
        z = r_coeff['r6'] * cols + r_coeff['r7'] * rows + r_coeff['r8']

        self.lut[:, :, 0] = x
        self.lut[:, :, 1] = y
        self.lut[:, :, 2] = z

        # debug
        # np.savetxt('/home/dataset_b/tmp_debug.txt', self.lut.reshape(-1, 3), fmt='%.6f', header='x y z', comments='')

    def align_dp_to_rgb(self, dp_img, t_coeff: dict):
        self.dp_extend.fill(0)

        for row in range(self.src_h):
            for col in range(self.src_w):
                dp_val = dp_img[row, col]
                x, y, z = self.lut[row, col]

                if dp_val > 0:
                    uv_color = [dp_val * x + t_coeff['t0'],
                                dp_val * y + t_coeff['t1'],
                                dp_val * z + t_coeff['t2']]
                    inv_z = 1.0 / uv_color[2]

                    new_x = int(uv_color[0] * inv_z + 0.5)
                    new_y = int(uv_color[1] * inv_z + 0.5)

                    if self.rotate_angle == 0:
                        pass
                    elif self.rotate_angle == 90:
                        new_x, new_y = new_y, new_x
                        new_y = self.dst_h - 1 - new_y
                    elif self.rotate_angle == 180:
                        new_x = self.dst_w - 1 - new_x
                        new_y = self.dst_h - 1 - new_y
                    elif self.rotate_angle == 270:
                        new_x, new_y = new_y, new_x
                        new_x = self.dst_w - 1 - new_x

                    if self.rotate_angle == 0:
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                    elif self.rotate_angle == 90:
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                    elif self.rotate_angle == 180:
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                    elif self.rotate_angle == 270:
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                    else:
                        pass

                else:
                    continue

        align_dp_img = self.dp_extend[self.extend_h_pixel:self.extend_buf_h - self.extend_h_pixel,
                                      self.extend_w_pixel:self.extend_buf_w - self.extend_w_pixel]

        return align_dp_img

    def align_dp_and_ir_to_rgb(self, ir_img, dp_img, t_coeff: dict):
        for row in range(self.src_h):
            for col in range(self.src_w):
                ir_val = ir_img[row, col]
                dp_val = dp_img[row, col]
                x, y, z = self.lut[row, col]

                if dp_val > 0:
                    uv_color = [dp_val * x + t_coeff['t0'],
                                dp_val * y + t_coeff['t1'],
                                dp_val * z + t_coeff['t2']]
                    inv_z = 1.0 / uv_color[2]

                    new_x = int(uv_color[0] * inv_z + 0.5)
                    new_y = int(uv_color[1] * inv_z + 0.5)

                    if self.rotate_angle == 0:
                        pass
                    elif self.rotate_angle == 90:
                        new_x, new_y = new_y, new_x
                        new_y = self.dst_h - 1 - new_y
                    elif self.rotate_angle == 180:
                        new_x = self.dst_w - 1 - new_x
                        new_y = self.dst_h - 1 - new_y
                    elif self.rotate_angle == 270:
                        new_x, new_y = new_y, new_x
                        new_x = self.dst_w - 1 - new_x

                    if self.rotate_angle == 0:
                        # process ir
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x + 1] = ir_val

                        # process dp
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                    elif self.rotate_angle == 90:
                        # process ir
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x + 1] = ir_val

                        # process dp
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x + 1] = int(uv_color[2])
                    elif self.rotate_angle == 180:
                        # process ir
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x - 1] = ir_val

                        # process dp
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y - 1, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                    elif self.rotate_angle == 270:
                        # process ir
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = ir_val
                        self.ir_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x - 1] = ir_val

                        # process dp
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x] = int(uv_color[2])
                        self.dp_extend[self.extend_h_pixel + new_y + 1, self.extend_w_pixel + new_x - 1] = int(uv_color[2])
                    else:
                        pass

                else:
                    continue

        align_dp_img = self.dp_extend[self.extend_h_pixel:self.extend_buf_h - self.extend_h_pixel,
                                      self.extend_w_pixel:self.extend_buf_w - self.extend_w_pixel]

        align_ir_img = self.ir_extend[self.extend_h_pixel:self.extend_buf_h - self.extend_h_pixel,
                                      self.extend_w_pixel:self.extend_buf_w - self.extend_w_pixel]

        return align_dp_img, align_ir_img
